import torch
from torch_geometric.data import DataLoader
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
# encoding=utf-8
from pylab import *  # 支持中文



class DataProcesser:
    @staticmethod
    def get_data_from_geometric_to_tensor(data):
        """
            获取数据并将geometric格式转换为tensor格式
        :param data: geometric格式数据
        :return: tensor格式数据
        """
        res_dict={}
        pass


    @staticmethod
    def get_data_loader(dataset,args):
        train_size = int(len(dataset) * args.training_set_proportion)
        test_size = len(dataset) - train_size
        train_dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])
        train_dataLoader = DataLoader(dataset=train_dataset, batch_size=args.batch_size,shuffle=False)
        test_dataLoader = DataLoader(dataset=test_dataset, batch_size=args.batch_size,shuffle=False)
        return train_dataLoader,test_dataLoader

    @staticmethod
    def graph_showing(data):
        '''
        args:
             data: torch_geometric.data.Data
        '''
        G = nx.Graph()
        x,edge_index,edge_attr,y=data
        # for ea in edge_attr[1]:
        #     pass
        for ei,ea in zip(edge_index[1].t(),edge_attr[1]):
            G.add_edge(ei[0], ei[1], name=ea, weight=50)
        for xx in x[1]:
            G.add_node(xx)
        for yy in y[1]:
            pass
        # topology construction logic
        # G.add_node('s1', desc='I am switch 1', attr1='new attr1')
        # G.add_node('s2', desc='I am switch 2', attr1='new attr1')
        # G.add_node('s3', desc='I am switch 3', attr1='new attr1')
        # G.add_edge('s1', 's2', name='edge 0', weight=50)
        # G.add_edge('s1', 's3', name='edge 1', weight=100)
        # G.add_edge('s2', 's3', name='edge 2', weight=20)

        # draw graph with labels
        pos = nx.spring_layout(G)
        nx.draw(G, pos)

        # generate node_labels manually
        node_labels = {}
        for node in G.nodes:
            node_labels[node] = G.nodes[node]  # G.nodes[node] will return all attributes of node

        nx.draw_networkx_labels(G, pos, labels=node_labels)

        # generate edge_labels manually
        edge_labels = {}
        for edge in G.edges:
            edge_labels[edge] = G[edge[0]][edge[1]]  # G[edge[0]][edge[1]] will return all attributes of edge

        nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)

        plt.show()

    @staticmethod
    def Line_chart_generator(x,y,title="My_model - loss",label=u'loss of training set'):
        mpl.rcParams['font.sans-serif'] = ['SimHei']
        plt.plot(x, y, marker='', mec='r', mfc='w',label=label)
        plt.legend()  # 让图例生效
        # plt.subplots_adjust(bottom=0.15)
        plt.xlabel(u"epoch")  # X轴标签
        plt.ylabel("loss") #Y轴标签
        plt.title(title) #标题
        plt.show()
